# author: 徐溪键
import socket
from utils.SystemInfo import SystemInfo
from utils.Vulnerability import Vulnerabilty
from utils.BaseLineCheck import BaseLineCheck
from utils.Type import Type
import time

class LocalClient(object):
    def __init__(self, host, port, username):
        self.__host = host
        self.__port = port
        self.__username = username

        # 系统基本信息获取
        self.__systemInfo = SystemInfo()
        self.__uuid = self.__systemInfo.get_uuid()

        # 影子账户 映像劫持检测
        self.__vulnerablity = Vulnerabilty()

        # 基线检查
        self.__base_line_check = BaseLineCheck()

        # 加固策略

        # 接口映射表
        self.__function_map = {

            # 系统基本信息
            Type.BASE_INFO: self.__systemInfo.get_base_info,
            # 网络信息
            Type.NETWORK_INFO: self.__systemInfo.get_network_info,
            # 软件列表
            Type.INSTALLED_SOFT: self.__systemInfo.get_install_soft,
            # 补丁信息
            Type.PATCH_INFO: self.__systemInfo.get_patch_info,
            # 端口信息
            Type.PORT_INFO: self.__systemInfo.get_port_info,
            # 影子账户
            Type.SHADOW_ACCOUNT: self.__vulnerablity.check_shadow_account,
            # 映像劫持
            Type.IMAGE_HIJACK: self.__vulnerablity.check_image_hijack,
            # 基线检查信息
            Type.BASELINE_INFO: self.__base_line_check.check,
            # 基线检查扫描结果
            Type.BASELINE_CHECK_RESULT_INFO: self.__base_line_check.get_bc_result
        }

    @staticmethod
    def __encode_send_data(meta_data, username, UUID, type_code, date):
        # 将需要发送的数据处理并同时编码为utf-8
        meta_data = meta_data.replace('"', '\'')
        # 特殊处理转义字符\
        meta_data = meta_data.replace('\\', '__')
        # 待传输数据
        send_data = '{"typeCode":%d,"username":"%s","uuid":"%s","date":"%s","data":"%s"}' % (
            type_code, username, UUID, date, meta_data)
        return send_data.encode('utf-8')

    def __scan_info(self, type_code,date):
        # 建立socket对象进行信息数据传输
        client_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        client_socket.connect((self.__host, self.__port))
        # 调用接口映射表
        scan_meta_data = self.__function_map[type_code]()
        send_data = self.__encode_send_data(
            scan_meta_data, self.__username,  self.__uuid, type_code, date)
        # 传输数据
        client_socket.send(send_data)
        client_socket.close()

    def scan_all_info(self):
        # 为扫描传输数据设立时间
        date = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
        # 系统基本信息扫描传输
        self.__scan_info(Type.BASE_INFO,date)
        # 网络基本信息传输
        self.__scan_info(Type.NETWORK_INFO,date)
        # 软件列表信息扫描传输
        self.__scan_info(Type.INSTALLED_SOFT,date)
        # 补丁信息扫描传输
        self.__scan_info(Type.PATCH_INFO,date)
        # 端口信息扫描传输
        self.__scan_info(Type.PORT_INFO,date)
        # 影子账户信息扫描传输
        self.__scan_info(Type.SHADOW_ACCOUNT,date)
        # 映像劫持信息扫描传输
        self.__scan_info(Type.IMAGE_HIJACK,date)
        # 基线检查信息扫描传输
        self.__scan_info(Type.BASELINE_INFO,date)
        # 基线检查扫描结果传输
        self.__scan_info(Type.BASELINE_CHECK_RESULT_INFO,date)


# def main():
#     l= LocalClient(1, 1)


# if __name__ == '__main__':
#     main()
